package org.zlb.algorithm.tree;

import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

public class HuffmanTreeBuild {

    /**
     * 根据字符串构建哈夫曼树
     * @param s
     * @return
     */
    public static HuffmanTree build(String s) {
        // 可优化，如果是英文字母，可以采用数组方式
        Map<String, Integer> map = new HashMap<>();
        char[] arr = s.toCharArray();
        for (int i=0; i<arr.length; i++) {
            char ch = arr[i];
            String key = String.valueOf(ch);
            int w = 0;
            if (map.containsKey(key)) {
                w = map.get(key);
            }
            w += 1;
            map.put(key, w);
        }
        
        Map<String, HuffmanTree> treeMap = new HashMap<>();
        for (Entry<String, Integer> entry : map.entrySet()) {
            HuffmanTree tree = new HuffmanTree(entry.getKey(), entry.getValue());
            treeMap.put(entry.getKey(), tree);
        }
        
        HuffmanTree root = null;
        while (!treeMap.isEmpty()) {
            HuffmanTree minTree = minWeight(treeMap);
            treeMap.remove(minTree.getName());
            if (treeMap.isEmpty()) {
                root = minTree;
                break;
            }
            
            HuffmanTree nextTree = minWeight(treeMap);
            treeMap.remove(nextTree.getName());
            
            root = merge(minTree, nextTree);
            treeMap.put(root.getName(), root);
        }
        
        clearBranch(root);
        
        return root;
    }
    
    /**
     * 从map中获取权重最小的一个
     * @param map
     * @return
     */
    private static HuffmanTree minWeight(Map<String, HuffmanTree> map) {
        if (map.isEmpty()) {
            return null;
        }
        
        int min = Integer.MAX_VALUE;
        HuffmanTree item = null;
        for (Entry<String, HuffmanTree> entry : map.entrySet()) {
            if (entry.getValue().getWeight() < min) {
                min = entry.getValue().getWeight();
                item = entry.getValue();
            }
        }
        
        return item;
    }
    
    /**
     * 合并2个哈夫曼树
     * @param tree1
     * @param tree2
     * @return
     */
    private static HuffmanTree merge(HuffmanTree tree1, HuffmanTree tree2) {
        if (tree1 == null)
            return tree2;
        if (tree2 == null)
            return tree1;
        
        HuffmanTree tree = new HuffmanTree(null, tree1.getWeight() + tree2.getWeight());
        String name = null;
        if (tree1.getWeight() < tree2.getWeight()) {
            tree.setLeft(tree1);
            tree.setRight(tree2);
            name = tree1.getName()+ "-" +tree2.getName();
        } else {
            tree.setLeft(tree2);
            tree.setRight(tree1);
            name = tree2.getName()+ "-" +tree1.getName();
        }
        tree.setName(name);
        
        return tree;
    }
    
    /**
     * 清理枝节点
     * @param tree
     */
    private static void clearBranch(HuffmanTree tree) {
        if (tree.getLeft() == null && tree.getRight() == null) {
            return;
        }
        
        if (tree.getLeft() != null) {
            clearBranch(tree.getLeft());
        }
        
        if (tree.getRight() != null) {
            clearBranch(tree.getRight());
        }
    }
    
}
